# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1--2.14
import sqlite3
import pandas as pd
from utils.instances import TOKENIZER, LLM
from utils.configFinRAG import sql_examples_path
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer

# Global lists to store example data
g_example_question_list = []
g_example_sql_list = []
g_example_fa_list = []
g_example_info_list = []
g_example_token_list = []


# Function to load example data once
def load_example_data():
    sql_examples_file = pd.read_csv(sql_examples_path, delimiter=",", header=0)

    # Storing the columns directly into the lists (no need for iteration)
    global g_example_question_list, g_example_sql_list, g_example_info_list, g_example_fa_list
    g_example_question_list = sql_examples_file['问题'].tolist()
    g_example_sql_list = sql_examples_file['SQL'].tolist()
    g_example_info_list = sql_examples_file['资料'].tolist()
    g_example_fa_list = sql_examples_file['FA'].tolist()

    # Tokenize all questions at once
    global g_example_token_list
    g_example_token_list = [TOKENIZER(q)['input_ids'] for q in g_example_question_list]


# Function to retrieve data from the database using SQL
def execute_sql_query(sql):
    # Using 'with' for automatic resource management
    with sqlite3.connect(
            '/Users/wanghr/Documents/八维研修/项目/999-项目/fay数字人/金融场景智能问答系统/bs_challenge_financial_14b_dataset/dataset/博金杯比赛数据.db') as conn:
        cs = conn.cursor()
        success_flag, exc_result = query_db(sql, cs)
    return exc_result


# Main function to generate answer based on the query
def sql_retrieve_chain(query):
    # Load example data only once
    if not g_example_question_list:
        load_example_data()

    # Generate the SQL query and prompt
    result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list, g_example_token_list)

    # Execute SQL query
    exc_result = execute_sql_query(sql)

    # Generate the final answer using the query result
    answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list, g_example_fa_list,
                             g_example_token_list)

    return answer
